import bpy
import numpy as np

# Define the dynamics of the Chen system
def chen_derivatives(state, a, b, c):
    x, y, z = state
    dx_dt = a * (y - x)
    dy_dt = -7 * x - x * z + b * y
    dz_dt = x * y - c * z
    return np.array([dx_dt, dy_dt, dz_dt])

# Fourth-order Runge-Kutta method
def runge_kutta_chen(state, dt, a, b, c):
    k1 = dt * chen_derivatives(state, a, b, c)
    k2 = dt * chen_derivatives(state + 0.5 * k1, a, b, c)
    k3 = dt * chen_derivatives(state + 0.5 * k2, a, b, c)
    k4 = dt * chen_derivatives(state + k3, a, b, c)
    return state + (k1 + 2 * k2 + 2 * k3 + k4) / 6

class CreateChenCurve(bpy.types.Operator):
    """Create a Chen system curve"""
    bl_idname = "object.create_chen_curve"
    bl_label = "Create Chen Curve"
    bl_options = {'REGISTER', 'UNDO', 'PRESET'}

    # 可调节参数
    t_start: bpy.props.FloatProperty(
        name="Time Start",
        description="Start time of the simulation",
        default=0.0
    )
    t_end: bpy.props.FloatProperty(
        name="Time End",
        description="End time of the simulation",
        default=10.0
    )
    sample_interval: bpy.props.FloatProperty(
        name="Sample Interval",
        description="Interval between sample points",
        default=0.01,
        min=0.0001
    )

    def execute(self, context):
        # Set initial conditions and parameters
        initial_state = np.array([7, 7, 22])

        # Parameters of the Chen system
        a = 35
        b = 28
        c = 3

        # Run the simulation
        state = initial_state.copy()
        points = []
        t = self.t_start
        while t < self.t_end:
            points.append(state.copy())
            next_state = runge_kutta_chen(state, self.sample_interval, a, b, c)
            state = next_state
            t += self.sample_interval

        # Create a new curve object
        curve_data = bpy.data.curves.new(name="ChenCurve", type='CURVE')
        curve_data.dimensions = '3D'
        polyline = curve_data.splines.new('POLY')
        polyline.points.add(len(points) - 1)

        # Set the points of the curve
        for i, point in enumerate(points):
            polyline.points[i].co = (*point, 1)

        # Create an object from the curve data
        curve_object = bpy.data.objects.new("ChenCurveObject", curve_data)
        bpy.context.collection.objects.link(curve_object)

        return {'FINISHED'}

    def draw(self, context):
        layout = self.layout
        layout.prop(self, "t_start")
        layout.prop(self, "t_end")
        layout.prop(self, "sample_interval")


def menu_func(self, context):
    self.layout.operator(CreateChenCurve.bl_idname)


def register():
    bpy.utils.register_class(CreateChenCurve)
    bpy.types.VIEW3D_MT_add.append(menu_func)


def unregister():
    bpy.utils.unregister_class(CreateChenCurve)
    bpy.types.VIEW3D_MT_add.remove(menu_func)


if __name__ == "__main__":
    register()